package udp

import (
	"fmt"
	"net"
	"smart-flow/container/config"
	"smart-flow/container/logger"
	"sync/atomic"
	"time"

	"go.uber.org/zap"
)

type ServerManager struct {
	servers          []*EchoUdpServer
	packetTotalCount int64
	byteTotalCount   int64
}

func (m *ServerManager) Start() bool {
	for i := 0; i < config.UdpPortCount; i++ {
		m.startOne(config.UdpPortStartPos + i)
	}
	if len(m.servers) < config.UdpPortCount/2 {
		m.Stop()
		return false
	}
	go m.Stat()
	return true
}

func (m *ServerManager) startOne(port int) {
	server := new(EchoUdpServer)
	if !server.Start(port) {
		logger.Logger.Error("[容器] udp 服务启动错误", zap.String("ip", config.IpAddr), zap.Int("port", port))
		return
	}
	m.servers = append(m.servers, server)
	config.AddUdpPort(port)
	logger.Logger.Info("[容器] udp 服务启动成功", zap.String("ip", config.IpAddr), zap.Int("port", port))
}

func (m *ServerManager) Stop() {
	for _, server := range m.servers {
		server.Stop()
	}
}

func (m *ServerManager) Stat() {
	ticker := time.NewTicker(config.UdpInterval * time.Second)
	defer ticker.Stop()
	for range ticker.C {
		for _, server := range m.servers {
			m.packetTotalCount += server.PacketCount()
			m.byteTotalCount += server.ByteCount()
		}
		if m.packetTotalCount == 0 || m.byteTotalCount == 0 {
			continue
		}
		logger.Logger.Info("[容器] udp 接收数据统计(KB)", zap.Int("interval", config.UdpInterval),
			zap.Int64("packetCount", m.packetTotalCount),
			zap.Int64("byteCount", m.byteTotalCount/config.KB))
		m.byteTotalCount = 0
	}
}

func (m *ServerManager) RunerCount() int {
	return len(m.servers)
}

type EchoUdpServer struct {
	conn        *net.UDPConn
	packetCount atomic.Int64
	byteCount   atomic.Int64
}

func (s *EchoUdpServer) Start(port int) bool {
	addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", config.IpAddr, port))
	if err != nil {
		logger.Logger.Error("[容器] udp 服务端口错误", zap.String("ip", config.IpAddr), zap.Int("port", port), zap.Error(err))
		return false
	}
	s.conn, err = net.ListenUDP("udp", addr)
	if err != nil {
		logger.Logger.Error("[容器] udp 服务监听错误", zap.String("ip", config.IpAddr), zap.Int("port", port), zap.Error(err))
		return false
	}
	if err = s.conn.SetReadBuffer(config.UdpReceiveBuffSize); err != nil { // 设置接收缓冲区大小
		logger.Logger.Error("[容器] udp 服务设置缓冲区错误", zap.String("ip", config.IpAddr), zap.Int("port", port), zap.Int("receiveBuffSize", config.UdpReceiveBuffSize), zap.Error(err))
		return false
	}
	go s.handleReceive()
	return true
}

func (s *EchoUdpServer) Stop() {
	if s.conn != nil {
		_ = s.conn.Close()
	}
}

func (s *EchoUdpServer) handleReceive() {
	defer s.Stop()
	buf := make([]byte, config.UdpReadBufferSize)
	for {
		n, _, err := s.conn.ReadFromUDP(buf)
		if err != nil {
			continue
		}
		// logger.Logger.Debug("[容器] udp 接收到数据", zap.String("ip", addr.IP.String()), zap.Int("port", addr.Port), zap.Int("len", n))
		s.packetCount.Add(1)
		s.byteCount.Add(int64(n))
	}
}

func (s *EchoUdpServer) PacketCount() int64 {
	return s.packetCount.And(0)
}

func (s *EchoUdpServer) ByteCount() int64 {
	return s.byteCount.And(0)
}
